'''
@author: morales
Updated Jan 13, 2010
'''
from PolymorphismTables import PolymorphismTables
from SiteFrequencySpectrum import *
import numpy.matlib as np
import scipy.stats.mstats as mstats
import math

class Polyspect:
	
	myOperator = None
	myParameters = None
	mySFSs = None
	
	myMean = 0.0
	myStdev = 1.0
	
	myRejectLower = -2.0
	myRejectUpper = 2.0
	
	myPolyspectValues = []

	
	def __init__(self, iLinOp, iParameters, iSFS=None):
		self.myOperator = iLinOp
		self.myParameters = iParameters
		self.mySFSs = iSFS

		self.myNumChromosomes = len(iLinOp)
		self.__prepare()
	
	def operateOnAll(self, SFSs):
		
		#print 'Average Seg Sites is ', self.getAverageSegSites( SFSs )
		
		ssList = []
		for sfsItr in SFSs:
			ss = self.operateOn(sfsItr)
			ssList.append(ss)
		
		return ssList
		
	def operateOn(self, SFS):
		ss = self.__operateOnSFS(SFS)
		return ss

	def operateOn_Normalize(self, SFS):
		ss = self.operateOn(SFS)
		normedSS = (ss - self.myMean) / self.myStdev
		return normedSS

	def getRejectionValues(self, alpha):
		ssList = self.operateOnAll( self.mySFSs)
		rValues = mstats.mquantiles( ssList, [alpha/2.0, 1 - alpha/2.0])
		
		self.myRejectLower = rValues[0]
		self.myRejectUpper = rValues[1]
		return rValues
	
	def getAverageSegSites(self):
		ssList = []
		
		for sfsItr in self.mySFSs:
			ssList.append( np.sum(sfsItr) )
		
		avg = mstats.tmean( ssList )
		return avg
		
	def testAltHypotheses(self, inList):
		rValue = None
		rValue = np.zeros( len(inList), int )
		
		for i, value in enumerate(inList):
			if self.isRejected(value):
				rValue[0,i] = 1
				
		return np.ravel(rValue)
		
	
	def isRejected(self, value):
		if (value < self.myRejectLower) or (value > self.myRejectUpper):
			return True
		return False

	def __prepare(self):
		
		if self.mySFSs is None:
			self.__generateSimulations()
					
		self.__operationOnAllSFSs()
		
		self.myMean = mstats.tmean(self.myPolyspectValues)
		self.myStdev = mstats.std(self.myPolyspectValues)
		self.getRejectionValues(0.05)
		
	def __generateSimulations(self):
		file_Null = './tmp/SFS_Null.ms'
		
		pt = PolymorphismTables()
		msCommand = self.myParameters.getMSString_Null( self.myParameters.getNumSimulations() )
		msOut = pt.runMS(msCommand, file_Null)
		pt.readMS(file_Null, msOut)
	
		self.mySFSs = SiteFrequencySpectrum(pt)
		return self.mySFSs

	def __operationOnAllSFSs(self):
		self.myPolyspectValues = []
		
		for sfsItr in self.mySFSs:
			ss = self.__operateOnSFS(sfsItr)
			self.myPolyspectValues.append(ss)
		
		return self.myPolyspectValues

	def __operateOnSFS(self, isfs):
		dp = self.myOperator * isfs
		ss = dp.sum()
		return ss
		
	def __str__(self):
		s = "Polyspect"
		
		for i in range(1, len(self.myOperator) ):
			s += '\t%1.3f' % (self.myOperator[i] )
		
		return s
				
	def getPolyspectDimension(self):
		return len(self.myOperator)
	
	def getPolyspectSegSites(self):
		return self.myParameters.getSegSites()
	
	def getPolyspectValues(self):
		return self.myPolyspectValues
	
	def getNullDemography(self):
		return str(self.myParameters.myDemography)
	
	def getMidPValue(self, value):
		
		d = abs(value - self.myMean)
		tally = 0

		for s in self.myPolyspectValues:
			dist = abs( s - self.myMean )
			if dist > d:
				tally += 1
		
		percent = 1.0 * tally / len(self.myPolyspectValues)
		return percent
	